import numpy as np


class ZVShaper:
    def __init__(self, omega_d, zeta, ts) -> None:
        super().__init__()

        self._omega_d = omega_d
        self._zeta = zeta
        self._ts = ts

    def shape(self, traj: np.ndarray) -> np.ndarray:
        K = np.exp((np.pi * self._zeta) / (np.sqrt(1 - np.power(self._zeta, 2))))
        a0 = K / (K + 1.0)
        a1 = 1 - a0

        delay_t = np.pi / self._omega_d
        delay_count = int(delay_t / self._ts)

        shaped_traj = []
        if len(traj.shape) == 2:
            shaped_traj = np.zeros((traj.shape[0] + delay_count, traj.shape[1]))
            for i in range(shaped_traj.shape[0]):
                if i < delay_count:
                    shaped_traj[i, :] = a0 * (traj[i, :] - traj[0, :]) + traj[0, :]
                elif i >= traj.shape[0]:
                    shaped_traj[i, :] = a0 * traj[-1, :] + a1 * traj[i - delay_count, :]
                else:
                    shaped_traj[i, :] = a0 * traj[i, :] + a1 * traj[i - delay_count, :]
        elif len(traj.shape) == 1:
            shaped_traj = np.zeros((traj.shape[0] + delay_count))
            for i in range(shaped_traj.shape[0]):
                if i < delay_count:
                    shaped_traj[i] = a0 * (traj[i] - traj[0]) + traj[0]
                elif i >= traj.shape[0]:
                    shaped_traj[i] = a0 * traj[-1] + a1 * traj[i - delay_count]
                else:
                    shaped_traj[i] = a0 * traj[i] + a1 * traj[i - delay_count]
        return shaped_traj
